import argparse
import json
import os
from typing import Iterable, List

import requests


# 打开并读取txt文件的全部内容
def get_prompt(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
        return content


def post_http_request(messages: List[dict], api_url: str,stop_token: List[str] = []) -> requests.Response:
    headers = {
        "User-Agent": "Test Client",
        "Authorization": "Bearer EMPTY",  # 添加Authorization头
        "Content-Type": "application/json"
    }
    pload = {
        "model": "Yi-1.5-9B-Chat",
        "messages": messages,
        "temperature": 0.7,
        "stream": False,
        "stop": stop_token,
        "skip_special_tokens": "false",
        "stop_token_ids": [7]
    }
    response = requests.post(api_url, headers=headers, json=pload, stream=True)
    return response

def get_response(response: requests.Response) -> List[str]:
    data = json.loads(response.content)
    output = [choice["message"]["content"] for choice in data["choices"]]
    return output


def generate_prompt(file_name:str):
    file_path = os.path.join('./split_story', file_name + '_split.txt')
    target_path = os.path.join('./prompt', file_name + '.txt')

    api_url = f"http://localhost:30014/v1/chat/completions"

    messages = [
        {"role": "system", "content": get_prompt(file_path)},
        {"role": "user", "content": get_prompt('./prompt/init_prompt.txt')}
    ]

    # stop_tokens = ["<|im_end|>", "<|/im_start|>"]
    stop_tokens = ["<|im_end|>", "<|im_start|>", "<|/im_start|>", "<|im_continue|>"]

    response = post_http_request(messages, api_url, stop_token=stop_tokens)
    output_str = ''.join(get_response(response))
    with open(target_path, 'w', encoding='utf-8') as f:
            f.write(output_str)

if __name__ == "__main__":
    generate_prompt('story1')
